/**
 * Copyright (c) 2010 Axa Holding Belgium, SA. All rights reserved.
 * This software is the confidential and proprietary information of the AXA Group.
 */
package chabernac.proxyserver;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.Socket;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.apache.log4j.Logger;

import chabernac.io.CachingOutputStream;
import chabernac.tools.HeaderTools;
import chabernac.tools.IOTools;

public class ProxyWorker {
  private static Logger LOGGER = Logger.getLogger( ProxyWorker.class );

  private Socket myClientSocket = null;

  private SocketFactory myProxySocketFactory = null;
  private PSocket myProxySocket = null;

  private InputStream myClientInputStream = null;
  private OutputStream myClientOutputStream = null;
  private InputStream myProxyInputStream = null;
  private OutputStream myProxyOutputStream = null;

  private iAuthenticationProvider myAuthenticationProvicer = null;

  private List<String> myHeaders = null;
  private CachingOutputStream myCachingOutputStream = null;

  private ExecutorService myStreamService = Executors.newFixedThreadPool( 2 );
  private int myFailedCounter = 0;

  private Object AUTHENTICATION_LOCK = new Object();


  public ProxyWorker ( Socket anClientSocket , SocketFactory  aProxySocketFactory , iAuthenticationProvider aAuthenticationProvider ) {
    super();
    myClientSocket = anClientSocket;
    myProxySocketFactory = aProxySocketFactory;
    myAuthenticationProvicer = aAuthenticationProvider;

  }

  public void run() throws IOException, InterruptedException{
    try{
      if(LOGGER.isDebugEnabled()) LOGGER.debug("Begin of ProxyWorker");
      initClientSocket();
//      resetProxySocket();

      CountDownLatch theLatch = new CountDownLatch(2);

      myStreamService.execute( new InputCopier(theLatch));
//      theStreamService.execute( new OutputCopier(theLatch));

      theLatch.await();

      if(LOGGER.isDebugEnabled()) LOGGER.debug( "End of worker");
    } finally {
      if(myClientSocket != null){
        myClientSocket.close();
      }
      if(myProxySocket != null){
        myProxySocket.close();
      }
    }
  }

  private void initClientSocket() throws IOException{
    myClientInputStream = myClientSocket.getInputStream();
    myClientOutputStream = myClientSocket.getOutputStream();
  }

  private synchronized void resetProxySocket() throws IOException{
    if(myProxySocket != null){
      myProxyInputStream.close();
      myProxyOutputStream.close();
      myProxySocket.close();
    }

    myProxySocket = myProxySocketFactory.createSocket( myHeaders );
    myProxyInputStream = myProxySocket.getInputStream();
    myProxyOutputStream = myProxySocket.getOutputStream();
  }

  private class InputCopier implements Runnable{
    private CountDownLatch myLatch = null;

    public InputCopier ( CountDownLatch aLatch ) {
      super();
      myLatch = aLatch;
    }

    public void run() {
      try {
        if(LOGGER.isDebugEnabled()) LOGGER.debug( "Begin of InputCopier");
        byte[] theHeaderBytes = IOTools.readUntill( myClientInputStream, "\r\n\r\n".getBytes() );

        myHeaders = extractHeaders( theHeaderBytes );


        //only now after caching the headers when can connect to the web server
        //if we do not use the proxy.
        resetProxySocket();

        //synchronize on authentication lock
        //this will cause the code to block if an authentication dialog is currently active
        //after the code unlocks the authentication header will be added to this request to the proxy
        synchronized(AUTHENTICATION_LOCK){
          //modify headers
          if(myProxySocket.isProxySocket()){
            examineAuthentication(myHeaders, true);
          }

          myStreamService.execute( new OutputCopier(myLatch));

          writeHeaders( myHeaders, myProxyOutputStream );

          //only if the content length is > 0 
          //copy remainders
          myCachingOutputStream = new CachingOutputStream(myProxyOutputStream);
          copyContent(myClientInputStream, myCachingOutputStream, myHeaders);

          myProxyOutputStream.flush();
          if(LOGGER.isDebugEnabled()) LOGGER.debug( "End of InputCopier ");
        }
      } catch ( Exception e ) {
        LOGGER.error( "An error occured while copying lines", e );
      } finally {
        myLatch.countDown();
      }
    }
  }

  private class SimpleInputCopier implements Runnable{
    private CountDownLatch myLatch = null;

    public SimpleInputCopier ( CountDownLatch aLatch ) {

      super();
      myLatch = aLatch;
    }

    public void run(){
      try {
        IOTools.copyStream( myClientInputStream, myProxyOutputStream );
      } catch ( IOException e ) {
        LOGGER.error("Error occured while copying streams", e);
      } finally {
        myLatch.countDown();
      }
    }
  }

  private class OutputCopier implements Runnable{
    private CountDownLatch myLatch = null;

    public OutputCopier ( CountDownLatch aLatch ) {
      super();
      myLatch = aLatch;
    }

    public void run() {
      try {
        if(LOGGER.isDebugEnabled()) LOGGER.debug( "Begin of OutputCopier");
//        while(true){
        byte[] theHeaderBytes = IOTools.readUntill( myProxyInputStream, "\r\n\r\n".getBytes() );

        List<String> theHeaders = extractHeaders( theHeaderBytes );

        synchronized ( AUTHENTICATION_LOCK ) {
          while(isFailedAuthentication( theHeaders)){
            //the authentication has failed
            //get the headers and bytes from the input
            examineAuthentication( myHeaders, false );

            //reset the connection to the proxy
            resetProxySocket();
            writeHeaders( myHeaders, myProxyOutputStream );
            myProxyOutputStream.write( myCachingOutputStream.getCache() );
            myProxyOutputStream.flush();

            theHeaderBytes = IOTools.readUntill( myProxyInputStream, "\r\n\r\n".getBytes() );

            theHeaders = extractHeaders( theHeaderBytes );
          }
        }

        myAuthenticationProvicer.successfullAuthentication();

        writeHeaders( theHeaders, myClientOutputStream );

        if(myHeaders.get( 0 ).startsWith( "CONNECT")){
          //at this point we must just connect the streams together and wait for the time to live time.
          if(LOGGER.isDebugEnabled()) LOGGER.debug( "Connect detected" );
          // the input copier will already be stopped, we create a new one that simply copies all data
          CountDownLatch theLatch = new CountDownLatch(1);
          new Thread(new SimpleInputCopier(theLatch)).start();
          IOTools.copyStream( myProxyInputStream, myClientOutputStream );
          theLatch.await();
        } else {
          //copy remainders
          copyContent(myProxyInputStream, myClientOutputStream, theHeaders);
        }


        myClientOutputStream.flush();
        if(LOGGER.isDebugEnabled()) LOGGER.debug( "End of OutputCopier");
//        }
      } catch ( Exception e ) {
        LOGGER.error( "An error occured while copying lines", e );
      } finally {
        myLatch.countDown();
      }
    }
  }

  private void copyContent(InputStream anInputStream, OutputStream anOutputStream, List<String> aHeaders) throws IOException{
    String theContentLength = HeaderTools.findHeader( aHeaders, "Content-Length" );
    String theTransferEncoding = HeaderTools.findHeader( aHeaders, "Transfer-Encoding" );
    if(theTransferEncoding != null && theTransferEncoding.equalsIgnoreCase( "chunked" )){
      //first read the length, this is hexadecimal
      IOTools.copyUntil( anInputStream, anOutputStream, "0\r\n\r\n".getBytes() );
    } else if(theContentLength != null){
      int theCLength = Integer.parseInt( theContentLength );
      if(theCLength > 0){
        IOTools.copyStream( anInputStream, anOutputStream, theCLength);
      }
    }
  }

  private void writeHeaders(List<String> aHeaders, OutputStream anOutputStream) throws IOException{
    for(String theHeader  : aHeaders){
//      System.out.println(theHeader);
      anOutputStream.write( theHeader.getBytes() );
      anOutputStream.write("\r\n".getBytes());
    }

    anOutputStream.write("\r\n".getBytes());
  }



  private List<String> extractHeaders(byte[] aBytes) throws IOException{
    BufferedReader theReader = new BufferedReader(new InputStreamReader(new ByteArrayInputStream(aBytes)));
    List<String> theHeaders = new ArrayList< String >();
    String theLine = null;
    while((theLine = theReader.readLine()) != null && !theLine.equals( "" )){
      theHeaders.add( theLine );
    }
    return theHeaders;
  }

  private void examineAuthentication( List< String > anHeaders, boolean isUseCached ) {
    for(Iterator< String > i = anHeaders.iterator();i.hasNext();){
      String theHeader = i.next();
      if(theHeader.startsWith( "Proxy-Authorization" )){
        i.remove();
      }
    }

    if(isUseCached){
      if(myAuthenticationProvicer.getCachedAuthentication() != null){
        anHeaders.add( "Proxy-Authorization: "  + myAuthenticationProvicer.getCachedAuthentication());
      }
    } else {
      anHeaders.add( "Proxy-Authorization: "  + myAuthenticationProvicer.basicAuthentication());
    }
  }

  private boolean isFailedAuthentication( List< String > anHeaders ) {
    Pattern thePattern = Pattern.compile( "HTTP/1.1 (\\d*) (.*)" );
    for(String theHeader : anHeaders){
      Matcher theMatcher = thePattern.matcher( theHeader );
      if(theMatcher.matches()){
        int theReturnCode = Integer.parseInt(theMatcher.group( 1 ));
        //not authenticated
        if(theReturnCode == 401 || theReturnCode == 407){
          myAuthenticationProvicer.failedAuthentication();
          myFailedCounter++;
          return true;
        }
      }
    }
    return false;
  }
}
